"""
Base Learner Interface for Continual Learning

This interface defines the core contract for continual learning systems,
combining all the specialized interfaces into a cohesive learner interface.
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, Generic, TypeVar, Callable
import torch
import torch.nn as nn

from learners.interfaces.lifecycle import (
    LifecycleManager,
    TaskStateManager,
    ConfigurationManager,
)
from learners.interfaces.memory import MemoryManager, ClassMeanManager
from learners.interfaces.model import ModelManager, NetworkStateManager
from learners.interfaces.training import (
    TrainingManager,
    OptimizerFactory,
    SchedulerFactory,
    LossFactory,
)
from learners.interfaces.evaluation import (
    EvaluationManager,
    AccuracyCalculator,
    MetricsEvaluator,
)
from learners.interfaces.event import EventEmitter
from learners.interfaces.multi_gpu import MultiGPUManager, DeviceManager

# Type variable for model type
T = TypeVar("T", bound=nn.Module)


class BaseLearnerInterface(
    Generic[T],
    LifecycleManager,
    TaskStateManager,
    ConfigurationManager,
    MemoryManager,
    ClassMeanManager,
    ModelManager,
    NetworkStateManager,
    TrainingManager,
    OptimizerFactory,
    SchedulerFactory,
    LossFactory,
    EvaluationManager,
    AccuracyCalculator,
    MetricsEvaluator,
    EventEmitter,
    MultiGPUManager,
    DeviceManager,
    ABC,
):
    """
    Comprehensive interface for continual learning systems

    This interface combines all the specialized interfaces to provide
    a complete contract for continual learning implementations.
    """

    # Core properties that all learners must provide
    _network: T
    _old_network: T

    @abstractmethod
    def __init__(
        self, args: Dict[str, Any], data_manager: Any, model_func: Callable[..., T]
    ):
        """
        Initialize the learner

        Args:
            args: Configuration arguments
            data_manager: Data manager instance
            model_func: Function to create model instance
        """
        pass


class IncrementalLearner(BaseLearnerInterface[T]):
    """
    Specialized interface for incremental learning scenarios

    This interface extends the base learner interface with methods
    specific to incremental/class-incremental learning.
    """

    @abstractmethod
    def get_cls_range(self, task_id: int) -> tuple[int, int]:
        """
        Get class range for a specific task

        Args:
            task_id: Task identifier

        Returns:
            Tuple of (start_class, end_class)
        """
        pass

    @abstractmethod
    def update_fc_for_new_classes(self, nb_classes: int) -> None:
        """
        Update final classifier layer for new classes

        Args:
            nb_classes: Number of new classes to add
        """
        pass

    @abstractmethod
    def consolidate_knowledge(self) -> None:
        """
        Consolidate knowledge after learning new task

        This method should implement any knowledge distillation,
        regularization, or consolidation techniques.
        """
        pass


class TaskAwareLearner(IncrementalLearner[T]):
    """
    Interface for task-aware continual learning

    This interface is for learners that have access to task identity
    during both training and inference.
    """

    @abstractmethod
    def set_task_id(self, task_id: int) -> None:
        """Set current task ID for inference"""
        pass

    @abstractmethod
    def get_task_specific_parameters(self, task_id: int) -> Dict[str, Any]:
        """Get parameters specific to a task"""
        pass

    @abstractmethod
    def isolate_task_parameters(self, task_id: int) -> None:
        """Isolate parameters for specific task during training"""
        pass


class TaskAgnosticLearner(IncrementalLearner[T]):
    """
    Interface for task-agnostic continual learning

    This interface is for learners that do not have access to task identity
    during inference (and possibly training).
    """

    @abstractmethod
    def infer_task_id(self, inputs: Any) -> int:
        """
        Infer task ID from inputs

        Args:
            inputs: Input data

        Returns:
            Predicted task ID
        """
        pass

    @abstractmethod
    def learn_task_boundaries(self, data_stream: Any) -> list[int]:
        """
        Learn task boundaries from data stream

        Args:
            data_stream: Continuous data stream

        Returns:
            List of detected task boundary indices
        """
        pass


class MetaLearner(BaseLearnerInterface[T]):
    """
    Interface for meta-learning based continual learning

    This interface extends the base learner for methods that use
    meta-learning or learning-to-learn approaches.
    """

    @abstractmethod
    def meta_train(self, meta_train_tasks: list[Any]) -> None:
        """
        Meta-training phase

        Args:
            meta_train_tasks: List of meta-training tasks
        """
        pass

    @abstractmethod
    def meta_test(self, meta_test_task: Any) -> Dict[str, Any]:
        """
        Meta-testing phase

        Args:
            meta_test_task: Meta-testing task

        Returns:
            Meta-testing results
        """
        pass

    @abstractmethod
    def adapt_to_task(self, task_data: Any, num_steps: int) -> None:
        """
        Adapt to new task using few-shot learning

        Args:
            task_data: Data for the new task
            num_steps: Number of adaptation steps
        """
        pass


class RehearsalBasedLearner(IncrementalLearner[T]):
    """
    Interface for rehearsal-based continual learning methods

    This interface specializes the incremental learner for methods
    that use experience replay or rehearsal mechanisms.
    """

    @abstractmethod
    def replay_memory_samples(self, batch_size: int) -> Any:
        """
        Sample from replay memory

        Args:
            batch_size: Number of samples to retrieve

        Returns:
            Batch of memory samples
        """
        pass

    @abstractmethod
    def update_replay_buffer(self, new_data: Any, new_targets: Any) -> None:
        """
        Update replay buffer with new data

        Args:
            new_data: New training data
            new_targets: New training targets
        """
        pass

    @abstractmethod
    def compute_replay_loss(self, memory_batch: Any) -> torch.Tensor:
        """
        Compute loss on replay memory

        Args:
            memory_batch: Batch from replay memory

        Returns:
            Replay loss tensor
        """
        pass


class RegularizationBasedLearner(IncrementalLearner[T]):
    """
    Interface for regularization-based continual learning methods

    This interface specializes the incremental learner for methods
    that use regularization to prevent catastrophic forgetting.
    """

    @abstractmethod
    def compute_regularization_loss(self) -> torch.Tensor:
        """
        Compute regularization loss to prevent forgetting

        Returns:
            Regularization loss tensor
        """
        pass

    @abstractmethod
    def update_importance_weights(self) -> None:
        """Update importance weights for parameters"""
        pass

    @abstractmethod
    def consolidate_importance_weights(self) -> None:
        """Consolidate importance weights after task completion"""
        pass


class ArchitecturalLearner(IncrementalLearner[T]):
    """
    Interface for architectural approaches to continual learning

    This interface specializes the incremental learner for methods
    that modify or extend the model architecture during learning.
    """

    @abstractmethod
    def expand_network(self, expansion_config: Dict[str, Any]) -> None:
        """
        Expand network architecture for new task

        Args:
            expansion_config: Configuration for network expansion
        """
        pass

    @abstractmethod
    def prune_network(self, pruning_config: Dict[str, Any]) -> None:
        """
        Prune network to maintain efficiency

        Args:
            pruning_config: Configuration for network pruning
        """
        pass

    @abstractmethod
    def allocate_task_specific_modules(self, task_id: int) -> None:
        """
        Allocate task-specific modules

        Args:
            task_id: Task identifier
        """
        pass
